import torch
import pandas as pd
import numpy as np
import const
import importlib
import argparse


class Evaluator(object):

    def __init__(self, category_topk=(1, 3, 5), lm_relative_threshold=0.1):
        self.reset()
        self.category_topk = category_topk
        self.lm_relative_threshold = lm_relative_threshold

    def reset(self):
        self.category_accuracy = []
        self.attr_tp = 0
        self.attr_gt_p = 0
        self.attr_pred_p = 0
        self.lm_vis_tp = 0
        self.lm_vis_gt_p = 0
        self.lm_vis_pred_p = 0
        self.lm_pos_tp = 0
        self.lm_pos_gt_p = 0
        self.lm_pos_pred_p = 0

    def category_topk_accuracy(self, output, target):
        with torch.no_grad():
            maxk = max(self.category_topk)
            batch_size = target.size(0)

            _, pred = output.topk(maxk, 1, True, True)
            pred = pred.t()
            correct = pred.eq(target.view(1, -1).expand_as(pred))

            res = []
            for k in self.category_topk:
                correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
                res.append(correct_k.mul_(100 / batch_size))
            for i in range(len(res)):
                res[i] = res[i].cpu().numpy()[0] / 100

            self.category_accuracy.append(res)

    def attr_count(self, output, target):
        pred = torch.argmax(output, dim=1).cpu().numpy()
        target = target.cpu().numpy()
        # 去除标记全部为0的
        pred = pred[target.sum(axis=1) != 0]
        target = target[target.sum(axis=1) != 0]
        tp = ((pred == 1) & (target == 1)).sum(axis=1)
        fp = ((pred == 1) & (target == 0)).sum(axis=1)
        tn = ((pred == 0) & (target == 0)).sum(axis=1)
        fn = ((pred == 0) & (target == 1)).sum(axis=1)
        # tp, ground truth positive, predicted positive
        self.attr_tp += tp.sum()
        self.attr_gt_p += (tp + fn).sum()
        self.attr_pred_p += (tp + fp).sum()

    def landmark_count(self, output, sample):
        lm_dist = torch.sqrt(
            torch.sum(torch.pow((output['lm_pos_output'] - sample['landmark_pos_normalized']), 2), dim=2)
        ).cpu().detach().numpy()
        gt_lm_vis = sample['landmark_vis'].cpu().numpy()
        pred_lm_vis = torch.argmax(output['lm_vis_output'], dim=1).cpu().numpy()

        self.lm_vis_tp += ((gt_lm_vis == 1) & (pred_lm_vis == 1)).sum()
        self.lm_vis_gt_p += (gt_lm_vis == 1).sum()
        self.lm_vis_pred_p += (pred_lm_vis == 1).sum()
        self.lm_pos_tp += ((gt_lm_vis == 1) & (pred_lm_vis == 1) & (lm_dist < self.lm_relative_threshold)).sum()
        self.lm_pos_gt_p += (gt_lm_vis == 1).sum()
        self.lm_pos_pred_p += (pred_lm_vis == 1).sum()

    def add(self, output, sample):
        self.category_topk_accuracy(output['category_output'], sample['category_label'])
        self.attr_count(output['attr_output'], sample['attr'])
        self.landmark_count(output, sample)

    def evaluate(self):
        category_accuracy = np.array(self.category_accuracy).mean(axis=0)
        category_accuracy_topk = {}
        for i, top_n in enumerate(self.category_topk):
            category_accuracy_topk[top_n] = category_accuracy[i]
        attr_recall = float(self.attr_tp) / self.attr_gt_p
        attr_precision = float(self.attr_tp) / self.attr_pred_p
        lm_vis_recall = float(self.lm_vis_tp) / self.lm_vis_gt_p
        lm_vis_precision = float(self.lm_vis_tp) / self.lm_vis_pred_p
        lm_pos_recall = float(self.lm_pos_tp) / self.lm_pos_gt_p
        lm_pos_precision = float(self.lm_pos_tp) / self.lm_pos_pred_p

        return {
            'category_accuracy_topk': category_accuracy_topk,
            'attr_recall': attr_recall,
            'attr_precision': attr_precision,
            'lm_vis_recall': lm_vis_recall,
            'lm_vis_precision': lm_vis_precision,
            'lm_pos_recall': lm_pos_recall,
            'lm_pos_precision': lm_pos_precision,
        }

def parse_args_and_merge_const():
    parser = argparse.ArgumentParser()
    parser.add_argument('conf', default='', type=str)
    args = parser.parse_args()
    if args.conf != '':
        new_conf = importlib.import_module(args.conf)
        for key, value in new_conf.__dict__.items():
            if not(key.startswith('_')):
                setattr(const, key, value)

